import math
import functools
import quimb as qu

from quimb.core import (make_immutable, eye as qeye, )

from utils import diag

@functools.lru_cache(maxsize=8)
def ptm_pauli(ixyz, **kwargs):
	ixyz_map = {0: 'i', 'i': 'i', 'I': 'i',
	            1: 'x', 'x': 'x', 'X': 'x',
	            2: 'y', 'y': 'y', 'Y': 'y',
	            3: 'z', 'z': 'z', 'Z': 'z'}
	ptm_op_map = {
		'i': lambda: qeye(4, **kwargs),
		'x': lambda: diag([1, 1, -1, -1], **kwargs),
		'y': lambda: diag([1, -1, 1, -1], **kwargs),
		'z': lambda: diag([1, -1, -1, 1], **kwargs),
	}
	ptm_op = ptm_op_map[ixyz_map.get(ixyz, 'i')]()
	make_immutable(ptm_op)
	return ptm_op

@functools.lru_cache(4)
def ptm_hadamard(dtype=complex, sparse=False):
	ptm_H = qu.qu([
		[1, 0, 0, 0],
		[0, 0, 0, 1],
		[0, 0, -1, 0],
		[0, 1, 0, 0],
	], dtype=dtype, sparse=sparse)
	make_immutable(ptm_H)
	return ptm_H

@functools.lru_cache(128)
def ptm_rotation(phi, xyz='x', dtype=complex, sparse=False, **kwargs):
	"""The PTM representation of single qubit rotation gate."""
	xyz_map = {1: 'x', 'x': 'x', 'X': 'x',
	           2: 'y', 'y': 'y', 'Y': 'y',
	           3: 'z', 'z': 'z', 'Z': 'z'}
	sin_phi, cos_phi = math.sin(phi), math.cos(phi)
	opmap = {
		'x': lambda: qu.qu([
			[1, 0, 0, 0],
			[0, 1, 0, 0],
			[0, 0, cos_phi, -sin_phi],
			[0, 0, sin_phi, cos_phi]
		], dtype=dtype, sparse=sparse, **kwargs),
		'y': lambda: qu.qu([
			[1, 0, 0, 0],
			[0, cos_phi, 0, sin_phi],
			[0, 0, 1, 0],
			[0, -sin_phi, 0, cos_phi]
		], dtype=dtype, sparse=sparse, **kwargs),
		'z': lambda: qu.qu([
			[1, 0, 0, 0],
			[0, cos_phi, -sin_phi, 0],
			[0, sin_phi, cos_phi, 0],
			[0, 0, 0, 1]
		], dtype=dtype, sparse=sparse, **kwargs),
	}
	ptm_R = opmap[xyz_map.get(xyz, 'x')]()
	make_immutable(ptm_R)
	return ptm_R

ptm_Rx = functools.partial(ptm_rotation, xyz='x')
ptm_Ry = functools.partial(ptm_rotation, xyz='y')
ptm_Rz = functools.partial(ptm_rotation, xyz='z')

@functools.lru_cache(4)
def ptm_T_gate(dtype=complex, sparse=False):
	"""The PTM of T-gate (pi/8 gate)."""
	return ptm_rotation(math.pi / 4, xyz="z", dtype=dtype, sparse=sparse)

@functools.lru_cache(4)
def ptm_S_gate(dtype=complex, sparse=False):
	"""The PTM of S-gate (phase gate)."""
	return ptm_rotation(math.pi / 2, xyz="z", dtype=dtype, sparse=sparse)

@functools.lru_cache(4)
def ptm_CNOT(dtype=complex, sparse=True):
	ptm_cnot = qu.qu([
		[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
		[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
		[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
		[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
		[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
	], dtype=dtype, sparse=sparse)
	make_immutable(ptm_cnot)
	return ptm_cnot
